# 引入所需库
import torch
from torchvision import datasets, transforms
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt

# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)

# 一、数据载入和预处理
# 使用torch加载MNIST数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# 提取特征，转为numpy数组
X_train = train_dataset.data.numpy().reshape(-1, 28*28)
y_train = train_dataset.targets.numpy()
X_test = test_dataset.data.numpy().reshape(-1, 28*28)
y_test = test_dataset.targets.numpy()

# 二分类任务：判断大小数（将标签分为0~4和5~9两类）
y_train_binary = (y_train >= 5).astype(int)
y_test_binary = (y_test >= 5).astype(int)

# 二、逻辑回归模型
# 1. 判断大小数
log_reg = LogisticRegression(max_iter=1000)
log_reg.fit(X_train, y_train_binary)
y_pred_binary = log_reg.predict(X_test)
print("逻辑回归二分类准确率：", accuracy_score(y_test_binary, y_pred_binary))

# 2. Softmax回归模型识别手写数字
softmax_reg = LogisticRegression(multi_class='multinomial', solver='lbfgs', max_iter=1000)
softmax_reg.fit(X_train, y_train)
y_pred_softmax = softmax_reg.predict(X_test)
print("Softmax回归多分类准确率：", accuracy_score(y_test, y_pred_softmax))

# 三、支持向量机
# 1. 判断大小数
svm_binary = SVC(kernel='linear')
svm_binary.fit(X_train, y_train_binary)
y_pred_svm_binary = svm_binary.predict(X_test)
print("SVM二分类准确率：", accuracy_score(y_test_binary, y_pred_svm_binary))

# 2. 一对多方法识别手写数字
svm_ovr = SVC(decision_function_shape='ovr')
svm_ovr.fit(X_train, y_train)
y_pred_ovr = svm_ovr.predict(X_test)
print("SVM一对多多分类准确率：", accuracy_score(y_test, y_pred_ovr))

# 3. 一对一方法识别手写数字
svm_ovo = SVC(decision_function_shape='ovo')
svm_ovo.fit(X_train, y_train)
y_pred_ovo = svm_ovo.predict(X_test)
print("SVM一对一多分类准确率：", accuracy_score(y_test, y_pred_ovo))

# 四、决策树
decision_tree = DecisionTreeClassifier()
decision_tree.fit(X_train, y_train)
y_pred_tree = decision_tree.predict(X_test)
print("决策树多分类准确率：", accuracy_score(y_test, y_pred_tree))

# 五、随机森林
for n_estimators in [10, 50, 100]:
    random_forest = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
    random_forest.fit(X_train, y_train)
    y_pred_rf = random_forest.predict(X_test)
    print(f"随机森林（{n_estimators}子树）多分类准确率：", accuracy_score(y_test, y_pred_rf))
